#include <stdio.h>
#include <stdlib.h>
#include <assert.h>
#include <math.h>
#include <mpi.h>
#include <pthread.h>
#include <papi.h>

typedef struct body body;
typedef struct node node;

#define SIZE 9999.0
#define THETA_SQUARE 0.25
#define SOFTEN 18
#define DAMPING 0.1
#define THREAD_NUM 16

double time_step;
unsigned int np, rank;

node* root;
body* list;
double num_ite;
unsigned int count;
double* accln_x,* accln_y;

pthread_barrier_t barr;

struct body {
  unsigned int num;

  unsigned int mass;

  double loc_x;
  double loc_y; 

  double vel_x;
  double vel_y;
  double accln_x;
  double accln_y;
};

struct node {
  node* children[4];
  node* parent;
  body* b[4];

  unsigned int mass;

  double x_start;
  double x_end;
  double y_start;
  double y_end;

  double loc_x;
  double loc_y;

  pthread_mutex_t mutex;
};

void inline * populate_list(FILE* ifile, unsigned int* count) {
  unsigned int i;
  fscanf(ifile, "%u", count);

  body* list = malloc(sizeof(body)*(*count));

  for(i=0; i<*count; i++) {
    list[i].num = i;
    fscanf(ifile, "%lf %lf %u %*u %*u", &list[i].loc_x, &list[i].loc_y, &list[i].mass);
    list[i].vel_x = 0;
    list[i].vel_y = 0;
    list[i].accln_x = 0;
    list[i].accln_y = 0;
  }

  return list;
}

void inline init_node(node* n, node* p, double x_start, double x_end, double y_start, double y_end) {
  n->children[0] = NULL;  n->children[1] = NULL;
  n->children[2] = NULL;  n->children[3] = NULL;
  n->parent = p;
  n->b[0] = NULL;  n->b[1] = NULL;
  n->b[2] = NULL;  n->b[3] = NULL;
  n->x_start = x_start;
  n->x_end = x_end;
  n->y_start = y_start;
  n->y_end = y_end;
  n->mass = 0;
  pthread_mutex_init(&n->mutex, NULL);
}

void inline place_body(node* n, body* b) {
  double x = b->loc_x;
  double y = b->loc_y;
  double x_mid = (n->x_start + n->x_end)/2;
  double y_mid = (n->y_start + n->y_end)/2;
  if (x_mid - x > 0) {
    if (y_mid - y > 0)
      n->b[0] = b;
    else
      n->b[1] = b;
  }
  else {
    if (y_mid - y > 0)
      n->b[2] = b;
    else
      n->b[3] = b;
  }
}

void build_tree_helper(node* n, body* b, unsigned int depth) {
  double x = b->loc_x;
  double y = b->loc_y;
  double x_mid = (n->x_start + n->x_end)/2;
  double y_mid = (n->y_start + n->y_end)/2;

  if (n->loc_x != 0 && n->mass != 0) {
    printf("%u %lf %u %lf %lf\n", b->mass, b->loc_x, n->mass, n->loc_x, n->x_start);
    exit(-1);
  }

  if (depth > 50) {
    pthread_mutex_lock(& n->mutex);
    n->loc_x = (b->mass*b->loc_x + n->mass*n->loc_x)/(b->mass + n->mass);
    n->loc_y = (b->mass*b->loc_y + n->mass*n->loc_y)/(b->mass + n->mass);
    n->mass += b->mass;
    pthread_mutex_unlock(& n->mutex);
    return;
  }
  depth++;
  //printf("%% %lf %lf || %lf %lf %lf || %lf %lf %lf\n", x, y, n->x_start, x_mid, n->x_end, n->y_start, y_mid, n->y_end);

  if (x_mid - x > 0) {
    if (y_mid - y > 0) {
      // quadrature 1
      pthread_mutex_lock(& n->mutex);
      if (n->b[0] == NULL && n->children[0] == NULL) {
        n->b[0] = b;
        pthread_mutex_unlock(& n->mutex);
        return;
      }

      if (n->b[0] != NULL && n->children[0] == NULL) {
        node* t = malloc(sizeof(node));
        init_node(t, n, n->x_start, x_mid, n->y_start, y_mid);
        place_body(t, n->b[0]);
        n->b[0] = NULL;
        n->children[0] = t;

        pthread_mutex_unlock(& n->mutex);
        return build_tree_helper(n->children[0], b, depth);
      }
      pthread_mutex_unlock(& n->mutex);

      return build_tree_helper(n->children[0], b, depth);
    }

    else {
      // quadrature 2
      pthread_mutex_lock(& n->mutex);
      if (n->b[1] == NULL && n->children[1] == NULL) {
        n->b[1] = b;
        pthread_mutex_unlock(& n->mutex);
        return;
      }

      if (n->b[1] != NULL && n->children[1] == NULL) {
        node* t = malloc(sizeof(node));
        init_node(t, n, n->x_start, x_mid, y_mid, n->y_end);
        place_body(t, n->b[1]);
        n->b[1] = NULL;
        n->children[1] = t;

        pthread_mutex_unlock(& n->mutex);
        return build_tree_helper(n->children[1], b, depth);
      }
      pthread_mutex_unlock(& n->mutex);

      return build_tree_helper(n->children[1], b, depth);
    }
  }
  else {
    if (y_mid - y > 0) {
      // quadrature 3
      pthread_mutex_lock(& n->mutex);
      if (n->b[2] == NULL && n->children[2] == NULL) {
        n->b[2] = b;
        pthread_mutex_unlock(& n->mutex);
        return;
      }

      if (n->b[2] != NULL && n->children[2] == NULL) {
        node* t = malloc(sizeof(node));
        init_node(t, n, x_mid, n->x_end, n->y_start, y_mid);
        place_body(t, n->b[2]);
        n->b[2] = NULL;
        n->children[2] = t;

        pthread_mutex_unlock(& n->mutex);
        return build_tree_helper(n->children[2], b, depth);
      }
      pthread_mutex_unlock(& n->mutex);

      return build_tree_helper(n->children[2], b, depth);
    }

    else {
      // quadrature 3
      pthread_mutex_lock(& (n->mutex));
      if (n->b[3] == NULL && n->children[3] == NULL) {
        n->b[3] = b;
        pthread_mutex_unlock(& n->mutex);
        return;
      }

      if (n->b[3] != NULL && n->children[3] == NULL) {
        node* t = malloc(sizeof(node));
        init_node(t, n, x_mid, n->x_end, y_mid, n->y_end);
        place_body(t, n->b[3]);
        n->b[3] = NULL;
        n->children[3] = t;

        pthread_mutex_unlock(& n->mutex);
        return build_tree_helper(n->children[3], b, depth);
      }
      pthread_mutex_unlock(& n->mutex);

      return build_tree_helper(n->children[3], b, depth);
    }
  }
}

void inline build_tree(node* root, body* list, unsigned int count, unsigned int id) {
  unsigned int i=0;
  unsigned int start = count/THREAD_NUM*id;
  unsigned int end;
  if (id+1 == THREAD_NUM)
    end = count;
  else
    end = start + count/THREAD_NUM;

  for(i=start; i<end; i++) {
    build_tree_helper(root, &(list[i]), 0);
  }
}

void sweep_up(node* root) {
  unsigned int i;

  for(i=0; i<4; i=i+1) {
    if (root->children[i] != NULL) {
      sweep_up(root->children[i]);

      root->loc_x = (root->mass * root->loc_x + root->children[i]->mass * root->children[i]->loc_x) / (root->mass + root->children[i]->mass);
      root->loc_y = (root->mass * root->loc_y + root->children[i]->mass * root->children[i]->loc_y) / (root->mass + root->children[i]->mass);
      root->mass = root->mass + root->children[i]->mass;
    }
    else if (root->b[i] != NULL) {
      root->loc_x = (root->mass * root->loc_x + root->b[i]->mass * root->b[i]->loc_x) / (root->mass + root->b[i]->mass);
      root->loc_y = (root->mass * root->loc_y + root->b[i]->mass * root->b[i]->loc_y) / (root->mass + root->b[i]->mass);
      root->mass = root->mass + root->b[i]->mass;
    }
  }
}

void inline cal_force_helper(node* n, body* b, unsigned int size) {
  double mag;
  double x_diff = (n->loc_x - b->loc_x);
  double y_diff = (n->loc_y - b->loc_y);
  double distance = x_diff*x_diff + y_diff*y_diff;

  double size_squre = size*size;

  if (size_squre/distance < THETA_SQUARE) {
    // use approximation
    distance += SOFTEN;
    distance = pow(distance, 1.5);
    mag = n->mass / distance;

    b->accln_x += mag * x_diff;
    b->accln_y += mag * y_diff;
    return;
  }

  else {
    unsigned int i=0;
    for(i; i<4; i++) {
      if (n->b[i] != NULL) {
        x_diff = (n->b[i]->loc_x - b->loc_x);
        y_diff = (n->b[i]->loc_y - b->loc_y);
        distance = x_diff*x_diff + y_diff*y_diff;
        distance += SOFTEN;
        distance = pow(distance, 1.5);
        mag = n->b[i]->mass / distance;

        b->accln_x += mag * x_diff;
        b->accln_y += mag * y_diff;
      }

      if (n->children[i] != NULL) {
        cal_force_helper(n->children[i], b, size/2);
      }
    }
  }

}

void inline time_advance(node* root) {
  unsigned int i=0;
  for(i; i<4; i++) {
    if (root->b[i] != NULL) {
      root->b[i]->vel_x += root->b[i]->accln_x*time_step;
      root->b[i]->vel_x *= DAMPING;
      root->b[i]->vel_y += root->b[i]->accln_y*time_step;
      root->b[i]->vel_y *= DAMPING;

      root->b[i]->loc_x += root->b[i]->vel_x * time_step;
      root->b[i]->loc_y += root->b[i]->vel_y * time_step;

      if (root->b[i]->loc_x > SIZE)
        root->b[i]->loc_x = SIZE;
      if (root->b[i]->loc_y > SIZE)
        root->b[i]->loc_y = SIZE;
      if (root->b[i]->loc_x < -SIZE)
        root->b[i]->loc_x = -SIZE;
      if (root->b[i]->loc_y < -SIZE)
        root->b[i]->loc_y = -SIZE;
    }

    if (root->children[i] != NULL)
      time_advance(root->children[i]);
  }
  pthread_mutex_destroy(&(root->mutex));
  free(root);
  return;
}

void time_advance_helper(node* root, int id) {
  if (id == 0) {
    // time advance
    unsigned int j=0;
    for(j=0; j<4; j++) {
      if (root->b[j] != NULL) {
        root->b[j]->vel_x += root->b[j]->accln_x*time_step;
        root->b[j]->vel_x *= DAMPING;
        root->b[j]->vel_y += root->b[j]->accln_y*time_step;
        root->b[j]->vel_y *= DAMPING;

        root->b[j]->loc_x += root->b[j]->vel_x * time_step;
        root->b[j]->loc_y += root->b[j]->vel_y * time_step;

        if (root->b[j]->loc_x > SIZE)
          root->b[j]->loc_x = SIZE;
        if (root->b[j]->loc_y > SIZE)
          root->b[j]->loc_y = SIZE;
        if (root->b[j]->loc_x < -SIZE)
          root->b[j]->loc_x = -SIZE;
        if (root->b[j]->loc_y < -SIZE)
          root->b[j]->loc_y = -SIZE;
      }
    }
  }

  if (root->children[id] != NULL) {
    time_advance(root->children[id]);
  }
}

void* thread_main(void* t) {
  unsigned int id = (unsigned int) t;
  unsigned int i=0;

  while(i < num_ite) {
    // create root for the tree
    if (id == 0) {
      root = malloc(sizeof(node));
      init_node(root, NULL, -SIZE, SIZE, -SIZE, SIZE);
    }
    pthread_barrier_wait(&barr);

    // build the tree using body list
    build_tree(root, list, count, id);
    pthread_barrier_wait(&barr);

    if (id == 0) {
      // calculate center loc and mass bottom up
      sweep_up(root);
    }
    pthread_barrier_wait(&barr);

    // calculate force for each body in the tree
    {
      unsigned int j, start, end, division;

      division = count/np;
      start = division*rank;
      if (rank+1 == np)
        end = count;
      else
        end = start+division;

      if (id == 0) {
        accln_x = malloc(sizeof(double)*count);
        accln_y = malloc(sizeof(double)*count);
      }
      pthread_barrier_wait(&barr);

      unsigned int thread_start, thread_end;
      thread_start = start + (end-start)/THREAD_NUM*id;
      if (id+1 != THREAD_NUM)
        thread_end = thread_start + (end-start)/THREAD_NUM;
      else
        thread_end = end;

      for(j=thread_start; j<thread_end; j++) {
        list[j].accln_x = 0;
        list[j].accln_y = 0;
        cal_force_helper(root, &(list[j]), SIZE);
        accln_x[j] = list[j].accln_x;
        accln_y[j] = list[j].accln_y;
      }
      pthread_barrier_wait(&barr);

      if (id == 0) {
        unsigned int diff, aggre_start, aggre_end;
        for(j=0; j<np; j++) {
          aggre_start = j*division;
          if (j+1 == np)
            aggre_end = count;
          else
            aggre_end = aggre_start+division;
          diff = aggre_end - aggre_start;

          MPI_Bcast(accln_x+aggre_start, diff, MPI_DOUBLE, j, MPI_COMM_WORLD);
          MPI_Bcast(accln_y+aggre_start, diff, MPI_DOUBLE, j, MPI_COMM_WORLD);
        }
      }

      pthread_barrier_wait(&barr);
      thread_start = start/THREAD_NUM*id;
      if (id+1 != THREAD_NUM)
        thread_end = thread_start + start/THREAD_NUM;
      else
        thread_end = start;
      for(j=thread_start; j<thread_end; j++) {
        list[j].accln_x = accln_x[j];
        list[j].accln_y = accln_y[j];
      }

      thread_start = end + (count-end)/THREAD_NUM*id;
      if (id+1 != THREAD_NUM)
        thread_end = thread_start + (count-end)/THREAD_NUM;
      else
        thread_end = count;
      for(j=thread_start; j<thread_end; j++) {
        list[j].accln_x = accln_x[j];
        list[j].accln_y = accln_y[j];
      }

      pthread_barrier_wait(&barr);
      if (id == 0) {
        if (rank == 0) {
          for(j=0; j<count; j++) {
            //printf("%lf  ",accln_y[j]);
          }
          //printf("\n");
        }
        free(accln_x);
        free(accln_y);
      }
    }
    pthread_barrier_wait(&barr);

    if (id == 0) {
      // time advance
      unsigned int j=0;
      for(j=0; j<4; j++) {
        if (root->b[j] != NULL) {
          root->b[j]->vel_x += root->b[j]->accln_x*time_step;
          root->b[j]->vel_x *= DAMPING;
          root->b[j]->vel_y += root->b[j]->accln_y*time_step;
          root->b[j]->vel_y *= DAMPING;

          root->b[j]->loc_x += root->b[j]->vel_x * time_step;
          root->b[j]->loc_y += root->b[j]->vel_y * time_step;

          if (root->b[j]->loc_x > SIZE)
            root->b[j]->loc_x = SIZE;
          if (root->b[j]->loc_y > SIZE)
            root->b[j]->loc_y = SIZE;
          if (root->b[j]->loc_x < -SIZE)
            root->b[j]->loc_x = -SIZE;
          if (root->b[j]->loc_y < -SIZE)
            root->b[j]->loc_y = -SIZE;
        }
      }
    }

    pthread_barrier_wait(&barr);
    if (root->children[id/4] != NULL) {
      time_advance_helper(root->children[id/4], id%4);
    }

    pthread_barrier_wait(&barr);
    if (id == 0) {
      pthread_mutex_destroy(&(root->mutex));
      free(root);
    }

    if (id == 0 && rank == 0) {
      unsigned int j;
      for(j=0; j<count; j++) {
        //printf("%lf %lf %u %lf %lf\n", list[j].loc_x, list[j].loc_y, list[j].mass, list[j].vel_x, list[j].vel_y);
      }
    }

    i++;
  }
  return NULL;
}

int main(int argc, char* argv[]) {
  assert(argc == 5 && "need extra input parameters");

  unsigned long chkflg = PAPI_VER_CURRENT;
  num_ite = atof(argv[1]);
  time_step = atof(argv[2]);
  double time = 0;
  FILE* ifile = fopen(argv[3], "r");
  unsigned int i=0;
  unsigned long start, end;

  pthread_t thread[THREAD_NUM];
  pthread_attr_t attr;
  void* status;

  // MPI init
  MPI_Init(&argc, &argv);
  MPI_Comm_size(MPI_COMM_WORLD, &np);
  MPI_Comm_rank(MPI_COMM_WORLD, &rank);

  /*init papi*/
  PAPI_library_init(chkflg);
  if (chkflg != PAPI_VER_CURRENT) {
    printf("Error PAPI Library out of date\n");
    exit(1);
  }

  start = PAPI_get_real_usec();

  // read info about all input bodies
  list = populate_list(ifile, &count);
  fclose (ifile);

  //PTHREAD init
  pthread_attr_init(&attr);
  pthread_barrier_init(&barr, NULL, THREAD_NUM);

  for(i=0; i<THREAD_NUM; i=i+1) {
    pthread_create(&thread[i], &attr, thread_main, (void*) i);
  }

  for(i=0; i<THREAD_NUM; i=i+1) {
    pthread_join(thread[i], &status);
  }

  end = PAPI_get_real_usec();
  printf("rank: %u, time: %u\n", rank, end-start);

  if (rank == 0) {
    FILE* ofile = fopen(argv[4], "w");
    fprintf(ofile, "%u\n", count);
    for(i=0; i<count; i++) {
      fprintf(ofile, "%.2lf %.2lf %u %.2lf %.2lf\n", list[i].loc_x, list[i].loc_y, list[i].mass, list[i].vel_x, list[i].vel_y);
    }
    fclose(ofile);
  }

  MPI_Finalize(); 
  return 0;
}
